// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this
// file except in compliance with the License. You may obtain a copy of the
// License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

// Usage:
//
// protoc --plugin=protoc-gen-flare-rpc=path/to/v2_plugin --cpp_out=./
//   --flare_rpc_out=./ proto_file

#include <sstream>
#include <string>

#include "glog/logging.h"
#include "google/protobuf/compiler/code_generator.h"
#include "google/protobuf/compiler/cpp/cpp_helpers.h"
#include "google/protobuf/compiler/plugin.h"
#include "google/protobuf/compiler/plugin.pb.h"
#include "google/protobuf/io/coded_stream.h"
#include "google/protobuf/io/zero_copy_stream.h"

#include "flare/base/encoding/hex.h"
#include "flare/base/logging.h"
#include "flare/base/string.h"
#include "flare/rpc/protocol/protobuf/plugin/async_decl_generator.h"
#include "flare/rpc/protocol/protobuf/plugin/basic_decl_generator.h"
#include "flare/rpc/protocol/protobuf/plugin/names.h"
#include "flare/rpc/protocol/protobuf/plugin/sync_decl_generator.h"

using namespace std::literals;

namespace flare::protobuf::plugin {

namespace {

std::string GetFilePrefix(
    const google::protobuf::FileDescriptor* file_descriptor) {
  constexpr auto kSuffix = ".proto"sv;
  auto prefix = file_descriptor->name();
  if (EndsWith(prefix, kSuffix)) {
    prefix.erase(prefix.end() - kSuffix.size(), prefix.end());
  }
  return prefix;
}

class V2CodeWriter : public CodeWriter {
 public:
  V2CodeWriter(const google::protobuf::FileDescriptor* file_descriptor,
               google::protobuf::compiler::GeneratorContext* generator_context)
      : generator_context_(generator_context) {
    filename_prefix_ = GetFilePrefix(file_descriptor);
    namespace_ = Replace(file_descriptor->package(), ".", "::");
  }

  std::string* NewInsertionToHeader(
      const std::string& insertion_point) override {
    return &hdr_insertions_[insertion_point].emplace_back();
  }

  std::string* NewInsertionToSource(
      const std::string& insertion_point) override {
    return &src_insertions_[insertion_point].emplace_back();
  }

  void Flush() {
    // Header.
    auto hdr_include_guard = "FLARE_PROTOBUF_PLUGIN_GENERATED_" +
                             EncodeHex(filename_prefix_) + "_H_";
    auto hdr = OpenWithExtension(".flare.pb.h");
    google::protobuf::io::CodedOutputStream hdr_stream(hdr.get());
    hdr_stream.WriteString(Format(
        "// Generated by Flare's Protocol Buffers plugin, DO NOT EDIT IT!\n"
        "// Source: {}.proto\n"
        "\n"
        "#ifndef {}\n"
        "#define {}\n"
        "\n",
        filename_prefix_, hdr_include_guard, hdr_include_guard));
    DumpInsertionsToFile(hdr_insertions_, &hdr_stream);
    hdr_stream.WriteString(
        Format("#endif  // {}", hdr_include_guard, hdr_include_guard));

    // Source file.
    auto src = OpenWithExtension(".flare.pb.cc");
    google::protobuf::io::CodedOutputStream src_stream(src.get());
    DumpInsertionsToFile(src_insertions_, &src_stream);
  }

 private:
  using Insertions = std::unordered_map<std::string, std::vector<std::string>>;

  std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream> OpenWithExtension(
      const std::string& ext) {
    return std::unique_ptr<google::protobuf::io::ZeroCopyOutputStream>(
        generator_context_->Open(filename_prefix_ + ext));
  }

  void DumpInsertionsToFile(const Insertions& insertions,
                            google::protobuf::io::CodedOutputStream* os) {
    auto&& try_get = [&](auto&& point) -> std::string {
      if (auto iter = insertions.find(point); iter != insertions.end()) {
        return Join(iter->second, "");
      }
      return "";
    };

    std::stringstream ss;
    ss << try_get(kInsertionPointIncludes) << "\n";
    if (!namespace_.empty()) {
      ss << "namespace " << namespace_ << " {\n\n";
    }
    ss << try_get(kInsertionPointNamespaceScope) << "\n";
    if (!namespace_.empty()) {
      ss << "}\n\n";
    }
    ss << try_get(kInsertionPointGlobalScope) << "\n";

    os->WriteString(ss.str());
  }

 private:
  google::protobuf::compiler::GeneratorContext* generator_context_;

  std::string filename_prefix_;
  std::string namespace_;
  Insertions hdr_insertions_;
  Insertions src_insertions_;
};

// The generator.
//
// It's implementation is scattered in several `generate_xxx_impl.cc` since it's
// so long that putting them all in a single TU is infeasible.
class V2Generator : public google::protobuf::compiler::CodeGenerator {
 public:
  bool Generate(const google::protobuf::FileDescriptor* file,
                const std::string& parameter,
                google::protobuf::compiler::GeneratorContext* generator_context,
                std::string* error) const override;

 private:
  void GeneratePrologue(const google::protobuf::FileDescriptor* file,
                        CodeWriter* writer) const;

  void GenerateCodeFor(const google::protobuf::FileDescriptor* file,
                       const google::protobuf::ServiceDescriptor* service,
                       CodeWriter* writer) const;

  void GenerateEpilogue(const google::protobuf::FileDescriptor* file,
                        CodeWriter* writer) const;
};

bool V2Generator::Generate(
    const google::protobuf::FileDescriptor* file, const std::string& parameter,
    google::protobuf::compiler::GeneratorContext* generator_context,
    std::string* error) const {
  if (!file->service_count()) {
    *error = Format("No service is defined in file [{}].", file->name());
    return false;
  }

  // We'll be generating code into this object. (See below for its functionality
  // and usage.)
  google::protobuf::compiler::CodeGeneratorResponse response;
  V2CodeWriter writer(file, generator_context);

  // Generate code.
  GeneratePrologue(file, &writer);
  for (int i = 0; i != file->service_count(); ++i) {
    GenerateCodeFor(file, file->service(i), &writer);
  }
  GenerateEpilogue(file, &writer);

  // We don't return anything to Protocol Buffers' runtime, instead, we wrote
  // everything we need into separate file via `writer`.
  writer.Flush();
  return true;
}

void V2Generator::GeneratePrologue(const google::protobuf::FileDescriptor* file,
                                   CodeWriter* writer) const {
  // Headers.
  auto header_incls = Format(
      "#include <utility>\n"
      "\n"
      // `.pb.h` comes before Protocol Buffers' header to avoid missing deps.
      "#include \"{}\"\n"
      "\n"
      "#include <google/protobuf/generated_enum_reflection.h>\n"
      "#include <google/protobuf/service.h>\n"
      "\n"
      "#include \"flare/base/callback.h\"\n"
      "#include \"flare/base/future.h\"\n"
      "#include \"flare/base/status.h\"\n"
      "#include \"flare/base/down_cast.h\"\n"
      "#include \"flare/base/maybe_owning.h\"\n"
      "#include \"flare/rpc/internal/stream.h\"\n",
      GetFilePrefix(file) + ".pb.h");
  auto source_incls =
      Format("#include \"{}\"\n", GetFilePrefix(file) + ".flare.pb.h") +
      "\n"
      "#include <mutex>\n" +
      "\n"
      "#include \"flare/rpc/rpc_channel.h\"\n"
      "#include \"flare/rpc/rpc_client_controller.h\"\n"
      "#include \"flare/rpc/rpc_server_controller.h\"\n";
  *writer->NewInsertionToHeader(kInsertionPointIncludes) = header_incls;
  *writer->NewInsertionToSource(kInsertionPointIncludes) = source_incls;

  // `RpcServerController` / `RpcClientController` brings in too much
  // dependencies, so we forward declare them so as not to bring them into the
  // header.
  *writer->NewInsertionToHeader(kInsertionPointIncludes) =
      "\n"
      "namespace flare {\n"
      "\n"
      "class RpcServerController;\n"
      "class RpcClientController;\n"
      "\n"
      "}  // namespace flare\n"
      "\n";

  // Initialize service descriptors. Indexed by service's indices.
  *writer->NewInsertionToSource(kInsertionPointNamespaceScope) = Format(
      "namespace {{\n"
      "namespace flare_rpc {{\n"
      "\n"
      "const ::google::protobuf::ServiceDescriptor*\n"
      "  file_level_service_descriptors[{service_count}];"
      "\n"
      "void InitServiceDescriptorsOnce() {{\n"
      "  static std::once_flag f;\n"
      "  std::call_once(f, [] {{\n"
      "    auto file = "
      "::google::protobuf::DescriptorPool::generated_pool()\n"
      "        ->FindFileByName(\"{file}\");\n"
      "    for (int i = 0; i != file->service_count(); ++i) {{\n"
      "      file_level_service_descriptors[i] = file->service(i);\n"
      "    }}\n"
      "  }});\n"
      "}}\n"
      "\n"
      "const ::google::protobuf::ServiceDescriptor*\n"
      "GetServiceDescriptor(int index) {{\n"
      "  InitServiceDescriptorsOnce();\n"
      "  return file_level_service_descriptors[index];\n"
      "}}\n"
      "\n"
      "}}  // namespace flare_rpc\n"
      "}}  // namespace\n"
      "\n",
      fmt::arg("service_count", file->service_count()),
      fmt::arg("file", file->name()));
}

void V2Generator::GenerateCodeFor(
    const google::protobuf::FileDescriptor* file,
    const google::protobuf::ServiceDescriptor* service,
    CodeWriter* writer) const {
  // Unlike generator V1, we don't have to bother generating class that provides
  // the same interface as `cc_generic_services`.

  // The synchronous one.
  SyncDeclGenerator().GenerateService(file, service, writer);
  SyncDeclGenerator().GenerateStub(file, service, writer);

  // The one uses future. We only generic stub here.
  //
  // TBH I don't think there will be any benefit in implementing server side as
  // `Async`-based (at least for now).
  //
  // Let's see if we'll need this in the future.
  AsyncDeclGenerator().GenerateStub(file, service, writer);

  if (!file->options().cc_generic_services()) {
    // If generic services is not being generated by `protoc`, we should do it.
    // This is necessary for rpc-mock to work.
    //
    // Note that, though, that we only need service (but not the stub) here.
    BasicDeclGenerator().GenerateService(file, service, writer);
    *writer->NewInsertionToHeader(kInsertionPointNamespaceScope) = Format(
        "using {} = {};\n"
        "\n",
        service->name(), GetBasicServiceName(service));
  }
}

void V2Generator::GenerateEpilogue(const google::protobuf::FileDescriptor* file,
                                   CodeWriter* writer) const {
  // Nothing yet.
}

}  // namespace

}  // namespace flare::protobuf::plugin

#ifndef FLARE_BUILD_WITH_BAZEL

// To recognize service option `gdt.streaming_response`, we need either:
//
// - Link with `//common/rpc:rpc_options_proto`, which we want to avoid in the
//   first place, or
//
// - Do some dirty hack.
struct StreamingResponseInitializer {
  StreamingResponseInitializer() {
    google::protobuf::internal::ExtensionSet::RegisterExtension(
        google::protobuf::MethodOptions::internal_default_instance(), 10003, 8,
        false, false);
  }
} streaming_response_initializer;

#endif

int main(int argc, char* argv[]) {
  // Make warning about using `gdt.stream_response` looks sane.
  google::InitGoogleLogging("");
  google::LogToStderr();
  FLAGS_log_prefix = false;

  flare::protobuf::plugin::V2Generator gen;
  return google::protobuf::compiler::PluginMain(argc, argv, &gen);
}
